package com.lizemin.decorate.filter.wrapper;

import com.lizemin.decorate.util.JSON;
import jakarta.servlet.ReadListener;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpMethod;
import org.springframework.util.StringUtils;
import org.springframework.web.util.HtmlUtils;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Enumeration;

/**
 * @author lzm
 * @date 2025/10/6 12:25
 * @description 请求包装类，用于对请求入参中的脚本做转义处理
 */
public class XssRequestWrapper extends HttpServletRequestWrapper {

    private static final Logger log = LoggerFactory.getLogger(XssRequestWrapper.class);

    /**
     * 缓存请求体，防止重复读取造成报错
     */
    private byte[] requestBody;

    /**
     * 构造函数，将原始的请求参数做js脚本转义处理，并且将结果缓存起来
     *
     * @param request 原始的HttpServletRequest对象
     */
    public XssRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        if (HttpMethod.GET.name().equalsIgnoreCase(request.getMethod())) {
            return;
        }
        ServletInputStream inputStream = super.getInputStream();
        byte[] bytes = inputStream.readAllBytes();
        if (bytes.length == 0) {
            return;
        }
        String content = new String(bytes);
        log.info("原始请求体为：{}", content);
        String escapedJson = JSON.getEscapedJson(content);
        log.info("转义后的请求体为：{}", escapedJson);
        this.requestBody = escapedJson.getBytes(StandardCharsets.UTF_8);
    }

    @Override
    public ServletInputStream getInputStream() {
        return getServletInputStream(this.requestBody);
    }

    /**
     * 将缓存中做过js脚本转义处理后的
     *
     * @param requestBody 缓存的请求体
     * @return
     */
    private static ServletInputStream getServletInputStream(byte[] requestBody) {
        return new ServletInputStream() {

            private int index = 0;

            @Override
            public int read() throws IOException {
                if (index >= requestBody.length) {
                    return -1;
                }
                return requestBody[index++] & 0xFF;
            }

            @Override
            public boolean isFinished() {
                return index >= requestBody.length;
            }

            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setReadListener(ReadListener readListener) {

            }
        };
    }

    /**
     * 获取请求体入参
     *
     * @return 会从缓存中读取js转义后请求入参的bufferReader
     */
    @Override
    public BufferedReader getReader() {
        return new BufferedReader(new InputStreamReader(getInputStream(), StandardCharsets.UTF_8));
    }

    /**
     * 获取url中的参数
     *
     * @param name 参数名称
     * @return 转义后的参数值
     */
    @Override
    public String getParameter(String name) {
        String parameter = super.getParameter(name);
        return HtmlUtils.htmlEscape(parameter);
    }

    /**
     * 获取url中的参数
     *
     * @param name 参数名称
     * @return 转义后的参数值
     */
    @Override
    public String[] getParameterValues(String name) {
        String[] parameterValues = super.getParameterValues(name);
        if (parameterValues != null) {
            for (int i = 0; i < parameterValues.length; i++) {
                parameterValues[i] = HtmlUtils.htmlEscape(parameterValues[i]);
            }
        }
        return parameterValues;
    }

    /**
     * 获取请求头中的参数
     *
     * @param name 参数名称
     * @return 转义后的参数值
     */
    @Override
    public Enumeration<String> getHeaders(String name) {
        Enumeration<String> headers = super.getHeaders(name);
        ArrayList<String> encodedHeaders = new ArrayList<>();
        headers.asIterator().forEachRemaining(header -> {
            encodedHeaders.add(HtmlUtils.htmlEscape(header));
        });
        return Collections.enumeration(encodedHeaders);
    }

    @Override
    public String getHeader(String name) {
        String header = super.getHeader(name);
        if (StringUtils.hasText(header)) {
            return HtmlUtils.htmlEscape(header);
        }
        return header;
    }

}
